from langchain_openai import ChatOpenAI
from langchain_core.tools import tool

from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.prebuilt import create_react_agent

from contextlib import ExitStack
from pydantic import BaseModel,Field
import requests
import json


class WeatherLoc(BaseModel):
    location:str=Field(description="The location name of the city")


@tool(args_schema=WeatherLoc)
def get_weather(location):
    """
        Function to query current weather.
        :param loc: Required parameter, of type string, representing the specific city name for the weather query. \
        Note that for cities in China, the corresponding English city name should be used. For example, to query the weather for Beijing, \
        the loc parameter should be input as 'Beijing'.
        :return: The result of the OpenWeather API query for current weather, with the specific URL request address being: https://api.openweathermap.org/data/2.5/weather. \
        The return type is a JSON-formatted object after parsing, represented as a string, containing all important weather information.
        """
    # Step 1.构建请求
    url = "https://api.openweathermap.org/data/2.5/weather"

    # Step 2.设置查询参数
    params = {
        "q": location,
        "appid": "5c939a7cc59eb8696f4cd77bf75c5a9a",  # 输入API key
        "units": "metric",  # 使用摄氏度而不是华氏度
        "lang": "zh_cn"  # 输出语言为简体中文
    }

    # Step 3.发送GET请求
    response = requests.get(url,params=params)

    # Step 4.解析响应
    data = response.json()
    return json.dumps(data)

tools = [get_weather]

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-4o",api_key=api_key,base_url=api_base)

stack = ExitStack()

checkpointer = stack.enter_context(SqliteSaver.from_conn_string(":memory:"))

graph = create_react_agent(llm,tools=tools,checkpointer=checkpointer)

config1 = {"configurable":{"thread_id":"222"}}

for chunk in graph.stream({"messages":["你好，我叫西瓜老师"]},config1,stream_mode="values"):
    chunk["messages"][-1].pretty_print()

for chunk in graph.stream({"messages":["请问我叫什么？"]},config1,stream_mode="values"):
    chunk["messages"][-1].pretty_print()

config2 = {"configurable": {"thread_id": "22"}}

for chunk in graph.stream({"messages": ["帮我查一下北京的天气"]}, config2, stream_mode="values"):
    chunk["messages"][-1].pretty_print()

for chunk in graph.stream({"messages": ["请问我刚才问了什么问题？"]}, config2, stream_mode="values"):
    chunk["messages"][-1].pretty_print()